from maix import camera, display, image, nn, app

detector = nn.YOLO11(model="/root/models/yolo11n.mud", dual_buff = True)

# 英文标签到中文标签的映射
label_map = {
    'person': '人',
    'bicycle': '自行车',
    'car': '汽车',
    'motorcycle': '摩托车',
    'airplane': '飞机',
    'bus': '公交车',
    'train': '火车',
    'truck': '卡车',
    'boat': '船',
    'traffic light': '交通灯',
    'fire hydrant': '消防栓',
    'stop sign': '停车标志',
    'parking meter': '停车计时器',
    'bench': '长椅',
    'bird': '鸟',
    'cat': '猫',
    'dog': '狗',
    'horse': '马',
    'sheep': '羊',
    'cow': '牛',
    'elephant': '大象',
    'bear': '熊',
    'zebra': '斑马',
    'giraffe': '长颈鹿',
    'backpack': '背包',
    'umbrella': '雨伞',
    'handbag': '手提包',
    'tie': '领带',
    'suitcase': '行李箱',
    'frisbee': '飞盘',
    'skis': '滑雪板',
    'snowboard': '滑雪板',
    'sports ball': '运动球',
    'kite': '风筝',
    'baseball bat': '棒球棒',
    'baseball glove': '棒球手套',
    'skateboard': '滑板',
    'surfboard': '冲浪板',
    'tennis racket': '网球拍',
    'bottle': '瓶子',
    'wine glass': '酒杯',
    'cup': '杯子',
    'fork': '叉子',
    'knife': '刀',
    'spoon': '勺子',
    'bowl': '碗',
    'banana': '香蕉',
    'apple': '苹果',
    'sandwich': '三明治',
    'orange': '橙子',
    'broccoli': '西兰花',
    'carrot': '胡萝卜',
    'hot dog': '热狗',
    'pizza': '披萨',
    'donut': '甜甜圈',
    'cake': '蛋糕',
    'chair': '椅子',
    'couch': '沙发',
    'potted plant': '盆栽植物',
    'bed': '床',
    'dining table': '餐桌',
    'toilet': '马桶',
    'tv': '电视',
    'laptop': '笔记本电脑',
    'mouse': '鼠标',
    'remote': '遥控器',
    'keyboard': '键盘',
    'cell phone': '手机',
    'microwave': '微波炉',
    'oven': '烤箱',
    'toaster': '烤面包机',
    'sink': '水槽',
    'refrigerator': '冰箱',
    'book': '书',
    'clock': '时钟',
    'vase': '花瓶',
    'scissors': '剪刀',
    'teddy bear': '泰迪熊',
    'hair drier': '吹风机',
    'toothbrush': '牙刷'
}

# 设置分辨率为1920*1080
cam = camera.Camera(1920, 1080, detector.input_format())
disp = display.Display()

# 图像宽度和高度
IMG_WIDTH = 1920
IMG_HEIGHT = 1080

while not app.need_exit():
    img = cam.read()
    objs = detector.detect(img, conf_th = 0.5, iou_th = 0.45)
    for obj in objs:
        img.draw_rect(obj.x, obj.y, obj.w, obj.h, color = image.COLOR_RED)
        # 获取英文标签并转换为中文标签
        english_label = detector.labels[obj.class_id]
        chinese_label = label_map.get(english_label, english_label)  # 如果没有找到映射，就使用英文标签
        msg = f'{chinese_label}: {obj.score:.2f}'
        img.draw_string(obj.x, obj.y, msg, color = image.COLOR_RED)
        
        # 计算物体中心点坐标
        center_x = obj.x + obj.w // 2
        center_y = obj.y + obj.h // 2
        
        # 将中心点坐标转换为百分比（精确到小数点后4位）
        center_x_percent = center_x / IMG_WIDTH
        center_y_percent = center_y / IMG_HEIGHT
        
        # 打印识别到的物体及其中心点（百分比形式）
        print(f'检测到: {chinese_label}, 置信度: {obj.score:.2f}, 中心点: ({center_x_percent:.4f}, {center_y_percent:.4f}), 边界框: ({obj.x}, {obj.y}, {obj.w}, {obj.h})')
    disp.show(img)